import json
import traceback

from tornado.websocket import WebSocketHandler
from tornado_request_mapping import request_mapping

from logger_config import logger
from services.digit_twin_service import DigitTwinService


@request_mapping("/ws/digit_twin")
class DigitTwinHandler(WebSocketHandler):
    client_set = set()
    cur_method_index = 0

    def __init__(self, application, request):
        super().__init__(application, request)
        self.service = DigitTwinService()

    def check_origin(self, origin):
        return True

    async def open(self):
        logger.info("WebSocket connection opened")
        self.client_set.add(self)
        try:
            data = await self.service.get_data()
            data = {
                'key': 'all',
                'data': data
            }
            await self.write_message(json.dumps(data, ensure_ascii=False))
        except Exception as e:
            logger.error(traceback.format_exc())

    def on_message(self, message):
        print(message)
        self.write_message(u"server said: " + message)

    def on_close(self):
        try:
            self.client_set.remove(self)
        except KeyError:
            pass
        # 前端如果关闭了websocket的页面，后端这里会自动调用
        logger.info("WebSocket closed")
